import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from lifelines import KaplanMeierFitter, CoxPHFitter
from lifelines.statistics import logrank_test, multivariate_logrank_test
import warnings

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

warnings.filterwarnings("ignore")


def gen_param():
    return 'day', 'i', 'cut'
    # return 'day', 'i', 'uncut'
    # return 'day', 'f', 'cut'
    # return 'day', 'f', 'uncut'
    # return 'hour', 'i', 'cut'
    # return 'hour', 'i', 'uncut'
    # return 'hour', 'f', 'cut'
    # return 'hour', 'f', 'uncut'
    # return 'minute', 'i', 'uncut'


# label_flag, int_flag, cut = gen_param()

categories = [
    'zd',
    'gender',
    'bmij',
    'nation',
    'j',
    'surgery',
    'chemotherapy',
    'targeted_drugs',
    'radiotherapy',
    'drugs',
    'food',
    'kenel',
    'fuzhang',
    'bsn',
    'gzy',
    'xr',
    'cdbb',
    'cdzb',
    'db',
    'bp',
    'ssfs',
    'fhfs',
    'icu',
    'bowelsound'
]

continues = [
    'age',
    'height',
    'weight',
    'index',
    'score',
    'time',
    'feces',
    'nrs_score',
    'changminyinyoushang',
    'changminyinzuoshang',
    'changminyinyouxia',
    'changminyinzuoxia',
    'wexner',
    'vaizey',
    'hb',
    'rbc',
    'wbc',
    'palb',
    'tp',
    'alb',
    'neut',
    'urea',
    'crea',
    'ua',
    'glu',
    'k',
    'na',
    'cl',
    'ca',
    'qcxr',
    'hr',
    'scxyssy',
    'scxyszy',
    'dxl',
    'mlpf',
    'xl',
    'xyss',
    'xysz',
    'mzsj',
    'cgsj',
    'qxsj',
    'sssj',
    'cxl',
    'xbl',
    'jtl',
    'jtry',
    'cdcd',
    'fqzs',
    'pqzs',
    'yy',
    'qcxl', 'qcxy', 'qcszy', 'spo', 'jcdxl', 'nrspf', 'ljsy', 'sy', 'ywjl', 'xhdb', 'hxb', 'bxb', 'qbdb', 'zdb', 'bdb',
    'zxlxb', 'cmz', 'nsz', 'jgz', 'ns', 'pttz', 'jia', 'naz', 'lz', 'gz', 'syzl', 'znl', 'ysl', 'cnznl', 'cnwznl',
    'ysxx', 'zsxx', 'yxxx', 'zxxx', 'jjmlpf',
]

data = pd.read_csv('../data/dataset-bia-day-i-cut-undummy.csv')
features = data.columns.tolist()
features.remove('label')
features.remove('event')


good = []
bad = []

for feature in features:
    fig, ax = plt.subplots()
    if feature in continues:
        data[feature] = data[feature].astype(float)
        data[feature] = data[feature].apply(lambda x: '1' if x > data[feature].median() else '0')
    for i in data[feature].unique():
        kmf = KaplanMeierFitter()
        df_tmp = data.loc[data[feature] == i]
        kmf.fit(df_tmp.label,
                event_observed=df_tmp.event,
                label=i)

        kmf.plot_survival_function(ci_show=False, ax=ax)

    p_value = multivariate_logrank_test(event_durations=data.label,
                                        groups=data[feature],
                                        event_observed=data.event).p_value

    if p_value <= 0.05:
        good.append(feature)
    else:
        bad.append(feature)

    p_value_text = ['p-value < 0.001' if p_value < 0.001 else 'p-value = %.4F' % p_value][0]
    ax.set_title("survival curves of {}\n logrank test {}".format(feature, p_value_text))

    # plt.savefig('km-{}-{}-{}-{}.png'.format(label_flag, int_flag, cut, feature))
    plt.savefig('../cox/png/{}.png'.format(feature))
    plt.show()

# pd.DataFrame(good).to_csv('km/good_feature.csv', index=None)
# pd.DataFrame(bad).to_csv('km/bad_feature.csv', index=None)
